"""
CWF (Channel-Wise Feature) 矩阵计算模块
基于皮尔逊相关系数计算每1秒EEG数据的通道交互特征矩阵

公式 (3.9): ρx,y = cov(X, Y) / (σx * σy) = E[(X - μx)(Y - μy)] / (σx * σy)
公式 (3.10): CWF_i,j = ρ(Xp,i, Xp,j)
"""
import numpy as np
from typing import Union, Tuple
import sys
import os

# 添加emotion目录到路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

class CWFCalculator:
    """
    CWF矩阵计算器
    专门计算每1秒EEG数据的通道交互特征矩阵
    """
    def __init__(self, n_channels: int = 32):
        """
        初始化CWF计算器
        Args:
            n_channels (int): EEG通道数，默认为32（DEAP数据集）
        """
        self.n_channels = n_channels
        
    def calculate_cwf_matrix(self, eeg_data: np.ndarray) -> np.ndarray:
        """
        计算CWF矩阵
        Args:
            eeg_data (np.ndarray): 1秒的EEG数据
                形状: (n_channels, n_samples) 或 (n_trials, n_channels, n_samples)     
        Returns:
            np.ndarray: CWF矩阵，形状为 (n_channels, n_channels)
        """
        # 处理不同维度的输入
        if eeg_data.ndim == 3:
            # (n_trials, n_channels, n_samples) -> 合并所有试次
            eeg_data = eeg_data.reshape(eeg_data.shape[0], -1).T  # (n_channels, n_trials*n_samples)
        elif eeg_data.ndim == 2:
            # (n_channels, n_samples) 或 (n_samples, n_channels)
            if eeg_data.shape[0] == self.n_channels:
                # (n_channels, n_samples)
                pass
            else:
                # (n_samples, n_channels) -> 转置
                eeg_data = eeg_data.T
        else:
            raise ValueError(f"不支持的输入维度: {eeg_data.ndim}，期望2D或3D")
        # 确保数据形状正确
        if eeg_data.shape[0] != self.n_channels:
            raise ValueError(f"通道数不匹配: 期望{self.n_channels}，实际{eeg_data.shape[0]}")
        # 使用numpy的corrcoef函数计算皮尔逊相关系数矩阵
        cwf_matrix = np.corrcoef(eeg_data)
        # 处理NaN值（当标准差为0时）
        cwf_matrix = np.nan_to_num(cwf_matrix, nan=0.0)
        # 确保对角线为1
        np.fill_diagonal(cwf_matrix, 1.0)
        return cwf_matrix
    
    def get_matrix_statistics(self, cwf_matrix: np.ndarray) -> dict:
        """
        获取CWF矩阵的统计信息
        Args:
            cwf_matrix (np.ndarray): CWF矩阵
        Returns:
            dict: 统计信息
        """
        # 提取上三角矩阵（排除对角线）
        upper_tri = cwf_matrix[np.triu_indices_from(cwf_matrix, k=1)]
        
        stats = {
            'mean_correlation': np.mean(upper_tri),
            'std_correlation': np.std(upper_tri),
            'min_correlation': np.min(upper_tri),
            'max_correlation': np.max(upper_tri),
            'strong_correlations': np.sum(np.abs(upper_tri) > 0.5),
            'weak_correlations': np.sum(np.abs(upper_tri) < 0.2)
        }
        
        return stats

def calculate_cwf_from_deap_data(subject_id: int = 1, data_dir: str = "emotion/DEAP_data_preprocessed"):
    """
    从DEAP数据计算CWF矩阵
    Args:
        subject_id (int): 被试ID，默认为1
        data_dir (str): 数据目录路径
    Returns:
        tuple: (eeg_matrices, cwf_matrices, labels)
    """
    print("=" * 60)
    print(f"从DEAP数据计算CWF矩阵 - 被试 {subject_id}")
    print("=" * 60)
    
    try:
        # 直接加载原始DEAP数据（不经过9x9矩阵预处理）
        print("正在加载原始DEAP数据...")

        # 导入DEAP加载器
        from deap_loader_simplified import DEAPLoader
        
        # 创建加载器
        loader = DEAPLoader(data_dir)
        
        # 加载原始数据
        subject_data = loader.load_subject(subject_id)
        
        # 提取原始EEG数据（前32个通道）
        full_eeg_data = loader.get_eeg_data(subject_data['data'])  # (40, 32, 7680)
        
        # 创建CWF计算器
        cwf_calc = CWFCalculator(n_channels=32)
    
        # 获取第一个试次的第一秒数据 (32, 128)
        first_trial_first_second = full_eeg_data[0, :, :128]  # 取前128个采样点（1秒）
   
        # 计算CWF矩阵
        cwf_matrix = cwf_calc.calculate_cwf_matrix(first_trial_first_second)

        # 获取统计信息
        stats = cwf_calc.get_matrix_statistics(cwf_matrix)
        
        print(f"\nCWF矩阵统计信息:")
        print(f"  平均相关系数: {stats['mean_correlation']:.6f}")
        print(f"  相关系数标准差: {stats['std_correlation']:.6f}")
        print(f"  最小相关系数: {stats['min_correlation']:.6f}")
        print(f"  最大相关系数: {stats['max_correlation']:.6f}")
        print(f"  强相关对数 (|r| > 0.5): {stats['strong_correlations']}")
        print(f"  弱相关对数 (|r| < 0.2): {stats['weak_correlations']}")
        
        # 获取标签
        labels = subject_data['labels']
        return first_trial_first_second, cwf_matrix, labels
        
    except FileNotFoundError as e:
        print(f"错误: {e}")
        print("请确保DEAP数据文件存在")
        return None, None, None
    except Exception as e:
        print(f"处理过程中出现错误: {e}")
        return None, None, None

if __name__ == "__main__":
    
    # 从DEAP数据计算CWF矩阵
    print("\n" + "=" * 60)
    print("从DEAP数据计算CWF矩阵")
    print("=" * 60)
    eeg_matrices, cwf_matrix, labels = calculate_cwf_from_deap_data(subject_id=1)
   